# _*_coding:utf-8_*_
"""
This time the module is for PIN test, only single 2410 is supported.
"""
from ATS.instrumentDriver import InstCom as inst
import time
import math

Func_Map = ["curr", "volt"]
debug_mode = 0
OK = 0
COMP_OCCURED = 101
### error code
MEAS_RNG_BIG = -101


def commonsetup(**kwargs):  # reset the specify instrument
    is_beep = kwargs.get("is_beep")
    is_rst = kwargs.get("is_rst")
    if is_rst:
        for s in is_rst:
            s_obj = inst.getinstobj(s)
            s_obj.reset()
    if is_beep:
        for s in is_beep:
            s_obj = inst.getinstobj(s)
            s_obj.setbeeper(is_beep[s])


def srcmeas(Sweep, Step, Bias, Open, Meas_Range, Src_Range, Compliance, Remote_Sense, Meas_Seq, Meas_Pro, Common):
    """
    For sweeping
    Single channel test. Use 2410 built in sweep-meas command
    Step: {}
    Bias: {}
    """
    print(('Common is ', Common))
    if Common["loop_mode"] == 1:
        outputonoff(Common, Sweep, Step, Bias, 1)  # set the output parameter
        rval = srcmeasloop(Sweep, Step, Bias, Meas_Seq, Meas_Pro, Common)
        outputonoff(Common, Sweep, Step, Bias, 0)
        for s in Meas_Pro:
            rval[s]["S"] = Common["data_pro"][s]
    else:
        if Common["gs_flag"] == 1:
            sweepmeastrigsetup(Sweep, Step, Bias, Meas_Seq, Meas_Pro, Src_Range, Common)
        else:
            samplemeastrigsetup(Bias, Meas_Seq, Meas_Pro, Src_Range, Common)
        outputonoff(Common, Sweep, Step, Bias, 1)
        starttrig(Common, Common["trig_smuch"])
        rval = getdatacheckcomp(Bias.keys() + Sweep.keys() + Step.keys(), Meas_Seq, Meas_Pro, Common, -1)
        outputonoff(Common, Sweep, Step, Bias, 0)
    # print("********",rval)
    if Common["is_returndata"] == 1:
        return rval
    else:
        return 0


def spotmeas(Meas_Seq, Common, result):
    for s in Meas_Seq:
        rval = universalcom({s: ["read?"]}, 3)
        if Meas_Seq[s][0] == 0 or Meas_Seq[s][0] == 2:
            result[s]["I"].append(rval[1])
        if Meas_Seq[s][0] == 1 or Meas_Seq[s][0] == 2:
            result[s]["V"].append(rval[0])
    return 0


def biassrc(Bias, Common, pts):
    """
        Set bias source
    """
    for s in Bias:
        s_obj = Common["smu_obj"][s]
        s_obj.setlevel(Bias[s][0], Bias[s][1])
        if Common.has_key("data_pro"):
            Common["data_pro"][s].extend([Bias[s][0]] * pts)


def checkcomp(smu_list, Common):
    rval = OK
    if Common["stop_comp"] == 0:
        return rval
    for s in smu_list:
        val = universalcom({s: [":stat:meas:cond?"]}, 2)
        # val = universalcom({s:[":SENS:CURR:PROT:TRIP?"]},2)
        code = eval(val)
        if code & 2 ** 14 == 2 ** 14 or code & 2 ** 7 == 2 ** 7:
            rval = COMP_OCCURED
            return rval
    return rval


def srcmeasloop(Sweep, Step, Bias, Meas_Seq, Meas_Pro, Common):
    # calc step value for sweep
    print("sweep is ", Sweep)
    for s in Sweep:
        s_param = Sweep[s][0]
        # get step val
        if s_param[0] == 0:  # linear mode
            s_param.append((s_param[-1] * 1.0 - s_param[-2]) / (Common["swp_pts"] - 1))
        elif s_param[0] == 1:  # myLog mode
            s_param.append((math.log10(abs(s_param[-1])) - math.log10(abs(s_param[-2]))) / (Common["swp_pts"] - 1))
    result = {}
    if Meas_Pro:
        Common["data_pro"] = {}
        for s in Meas_Pro:
            Common["data_pro"][s] = []
    SMUs = Sweep.keys() + Step.keys() + Bias.keys()
    print('Meas_Seq', Meas_Seq)
    for s in Meas_Seq:
        print(s)
        result[s] = {"V": [], "I": [], "T": [], "S": []}
    # t_start = time.clock()
    if Common["gs_flag"] == 0:
        biassrc(Bias, Common, Common["sample_pts"])
        for i in range(Common["sample_pts"]):
            # check compliance
            comp_occured = checkcomp(SMUs, Common)
            if comp_occured != COMP_OCCURED:
                if i == 0:
                    time.sleep(Common["hold_time"])
                time.sleep(Common["delay_time"])
            # t_post = time.clock()-t_start
            spotmeas(Meas_Seq, Common, result)
            if comp_occured == COMP_OCCURED:
                return result
    else:
        biassrc(Bias, Common, Common["swp_pts"] * Common["step_pts"])
        for i in range(Common["step_pts"]):
            tmp_bias = {}
            for s in Step:
                tmp_bias[s] = (Step[s][0][0] + i * Step[s][0][2], Step[s][1])
            biassrc(tmp_bias, Common, Common["swp_pts"])
            for j in range(Common["swp_pts"]):
                tmp_bias = {}
                for s in Sweep:
                    s_param = Sweep[s][0]
                    if s_param[0] == 0:
                        tmp_bias[s] = (s_param[-3] + j * s_param[-1], Sweep[s][1])
                    elif s_param[0] == 1:  # myLog mode
                        tmp_bias[s] = (s_param[-3] * 10 ** (s_param[-1] * j), Sweep[s][1])
                    else:  # list mode
                        tmp_bias[s] = (s_param[1 + j], Sweep[s][1])
                biassrc(tmp_bias, Common, 1)
                # check compliance
                comp_occured = checkcomp(SMUs, Common)
                if comp_occured != COMP_OCCURED:
                    if i == 0 and j == 0:
                        time.sleep(Common["hold_time"])
                    time.sleep(Common["delay_time"])
                # meas
                # t_post = time.clock()-t_start
                spotmeas(Meas_Seq, Common, result)
                if comp_occured == COMP_OCCURED:
                    return result
    # print("srcmeasloop",result)
    return result


def sweepmeastrigsetup(Sweep, Step, Bias, Meas_Seq, Meas_Pro, Src_Range, Common):
    """
    Single SMU for Sweep and Measure
    """
    src_smu = []
    # setup source
    for s in Sweep:
        src_smu.append(s)
        func = Func_Map[Sweep[s][1]]
        swp_param = Sweep[s][0]
        if swp_param[0] < 2:  # sweep
            lin_log = ("lin", "myLog")[swp_param[0]]
            # setup swp range
            cmd_str = ":sour:swe:rang %s" % ("auto", "fix", "best")[Src_Range[s][0]]
            Common["cmd_dict"][s].append(cmd_str)
            cmd_str = [":sour:%s:mode swe" % func, ":sour:swe:spac %s" % lin_log,
                       ":sour:%s:star %g" % (func, swp_param[1]),
                       ":sour:%s:stop %g" % (func, swp_param[2]),
                       ":sour:swe:poin %d" % Common["swp_pts"]]
            Common["cmd_dict"][s].extend(cmd_str)
        else:  # list mode
            # print("swp_param",swp_param)
            # {'SMU1': [[2, 0, 5.0, 10.0, 15.0, 20, 30, 47.5, 65.0, 82.5, 100], 1]}
            cmd_str = ":sour:%s:mode list;:sour:list:%s %s" % (func, func, str(swp_param[1:])[1:-1])
            Common["swp_pts"] = len(swp_param) - 1
            Common["cmd_dict"][s].append(cmd_str)
    # setup trigger
    for s in Sweep:
        # There is no cmd for arm layer delay control. Program will turn on the output at the startV before trigger model initialized.
        # arm source and trigger source is default by immediatly.
        cmd_str = ":trig:coun %d;:trig:del %g" % (Common["swp_pts"], Common["delay_time"])
        Common["cmd_dict"][s].append(cmd_str)
        Common["trig_smuch"] = {s: 1}
    # send cmd
    universalcom(Common["cmd_dict"])
    Common["est_t"] = Common["delay_time"] * Common["swp_pts"]
    Common["total_pts"] = Common["swp_pts"]
    # Not needed
    # for s in Sweep:#confirm setup executed
    #    cmd_str = "*opc?"
    #    Common["cmd_dict"][s].append(cmd_str)
    # universalcom(Common["cmd_dict"],2)


def samplemeastrigsetup(Bias, Meas_Seq, Meas_Pro, Src_Range, Common):
    """
    Single SMU for Sweep and Measure
    """
    src_smu = []
    # setup source
    for s in Bias:
        src_smu.append(s)
        func = Func_Map[Bias[s][1]]
        cmd_str = ":sour:%s:mode fix;:sour:%s %g" % (func, func, Bias[s][0])
        Common["cmd_dict"][s].append(cmd_str)
    # setup trigger
    for s in Bias:
        # There is no cmd for arm layer delay control. Program will turn on the output at the startV before trigger model initialized.
        # arm source and trigger source is default by immediatly.
        cmd_str = ":trig:coun %d;:trig:del %g" % (Common["sample_pts"], Common["delay_time"])
        Common["cmd_dict"][s].append(cmd_str)
        Common["trig_smuch"] = {s: 1}
    # send cmd
    universalcom(Common["cmd_dict"])
    Common["est_t"] = Common["delay_time"] * Common["sample_pts"]
    Common["total_pts"] = Common["sample_pts"]
    for s in Bias:  # confirm setup executed
        cmd_str = "*opc?"
        Common["cmd_dict"][s].append(cmd_str)
    universalcom(Common["cmd_dict"], 2)


def outputonoff(Common, Sweep, Step, Bias, on_off):
    # print("*********Common",Common)
    if on_off == 0:
        for s in Common["smu_info"]:
            s_obj = Common["smu_obj"][s]
            s_obj.setlevel(0, Common["smu_info"][s])
    if (on_off == 0 and Common.get("off_0V") != 1) or on_off == 1:
        for s in Bias:
            s_obj = Common["smu_obj"][s]
            s_obj.outputenable(on_off)
        for s in Sweep:
            s_obj = Common["smu_obj"][s]
            s_obj.outputenable(on_off)
    if on_off == 1:  # set level by start value for voltage settling
        for s in Sweep:
            swp_param = Sweep[s][0]
            s_obj = Common["smu_obj"][s]
            s_obj.setlevel(swp_param[1], Sweep[s][1])
        if Common.get("loop_mode") != 1:
            for s in Common["smu_info"]:
                cmd_str = ":syst:time:icons"
                Common["cmd_dict"][s].append(cmd_str)
    universalcom(Common["cmd_dict"])


def postdata(result, Common):
    if Common["is_postdata"] == 0: return
    for s in result:
        for var in result[s]:
            __postdata(var, result[s][var])


def __postdata(name, val_list):
    #    print("***",name, str(val_list))
    # to do
    pass
    # postarraydata(name, val_list)


def getdatacheckcomp(SMUs, Meas_Seq, Meas_Pro, Common, step_ind):
    result = {}
    # No real time data post this time
    if Common["delay_time"] > 0.05:
        get_int = Common["delay_time"]
    else:
        get_int = 0.05
    cmd_dict = {}
    smu_rcv = Meas_Seq.keys() + Meas_Pro.keys()
    for s in smu_rcv:
        # no support for stop-on-compliance check if "fetch?" used because data won't be got till sweep finished.
        # to support stop-on-compliance, "data?" instead
        cmd_dict[s] = ["fetch?"]
        result[s] = []
    universalcom(cmd_dict)
    print('Drive========cmd_dict', cmd_dict)
    time.sleep(Common["est_t"])
    while 1:
        end_flag = 1
        for s in smu_rcv:
            if len(result[s]) < Common["total_pts"]:
                stat = inst.spoll(s)
                # print("*****stat = ",stat)
                if stat & 16:  # read data
                    rval = universalcom({s: [""]}, 4)
                    # time.sleep(100)
                    result[s].extend(rval)
                end_flag = 0
        if end_flag == 1:
            break
        time.sleep(get_int)
    result_new = parsedata(result, Common)
    postdata(result_new, Common)

    for s in SMUs:  # confirm if test completed
        cmd_dict[s] = ["*opc?"]
    rval = universalcom(cmd_dict, 2)
    return result_new


def parsedata(result, Common):
    result_new = {}
    name_map = {"curr": 'I', "volt": 'V', "time": 'T'}
    for s in result:
        result_new[s] = {}
        div = len(Common["var_name"][s])
        length = len(result[s]) / len(Common["var_name"][s])
        for i, v in enumerate(Common["var_name"][s]):
            name = name_map[v]
            result_new[s][name] = []
            for ind in range(length):
                result_new[s][name].append(result[s][ind * div + i])
    return result_new


def starttrig(Common, trig_smuch):
    """
    cmd_str = ":trac1:cle;:trac2:cle;:trac:poin 100;:trac2:poin 100;:trac1:feed sens;:trac2:feed sens;:trac1:feed:cont next;:trac2:feed:cont next"
    inst.write("SMU1", cmd_str)
    print(cmd_str)
    rval = inst.ask("SMU1", ":SYST:ERR:ALL?")
    print("error:", rval)
    """
    time.sleep(Common["hold_time"])
    for s in trig_smuch:
        cmd_str = ":init"
        Common["cmd_dict"][s].append(cmd_str)
    universalcom(Common["cmd_dict"])


def _initcmddict(Open, Src_Range, Common):
    smu_info = {}
    cmd_dict = {}
    smu_obj = {}
    for s in Open:
        smu_info[s] = -1
        cmd_dict[s] = []
        smu_obj[s] = inst.getinstobj(s)
    for s in Src_Range:
        smu_info[s] = Src_Range[s][2]
        cmd_dict[s] = []
        smu_obj[s] = inst.getinstobj(s)
    # smu_info = {'SMU2': 1, 'SMU1': 1}
    # cmd_dict = {"SMU1":[],"SMU2":[]}
    Common["smu_info"] = smu_info
    Common["cmd_dict"] = cmd_dict
    Common["smu_obj"] = smu_obj

    # for s in smu_info:
    #    inst.write(s, "*rst")


def setup(Bias, Step, Sweep, Open, Meas_Range, Src_Range, Compliance, Remote_Sense, Meas_Seq, Meas_Pro, Common):
    """
    Step: {} not supported
    Open: {} not supported
    """
    _initcmddict(Open, Src_Range, Common)
    srcrng_over20 = 0
    for s in Src_Range:
        s_obj = Common["smu_obj"][s]
        s_obj.setsrcfunc(Src_Range[s][2])
        if Src_Range[s][0] == 1:  # fixed range
            if Src_Range[s][2] == 1 and Src_Range[s][1] > 20:  # To avoid "exceed power limit" error, if V range on 1KV
                setting = s_obj.getinstsetting(["limiti", "rangei"])
                if setting.get("limiti") > 0.02:
                    s_obj.limiti(0.02)
                if setting.get("rangei") > 0.02:
                    s_obj.rangei(0.02)
                srcrng_over20 = 1
            s_obj.srcrange(Src_Range[s][1], Src_Range[s][2])
        elif Src_Range[s][0] == 2:  # auto range with low limit
            s_obj.srclorange(Src_Range[s][1], Src_Range[s][2])
        else:  # auto range
            s_obj.srcrange(0, Src_Range[s][2])
    for s in Compliance:
        s_obj = Common["smu_obj"][s]
        s_obj.limit(Compliance[s][0], Compliance[s][1])
    Common["var_name"] = {}
    for s in Meas_Seq:
        s_obj = Common["smu_obj"][s]
        func_id = Meas_Seq[s][0]
        s_obj.setmeasfunc(func_id)
        if s not in Meas_Pro:
            func_str = ["curr", "volt", "volt,curr"][func_id] + ["", ",time"][Common["time_stamp"]]
            if func_id > 1:
                form_elem = [0, 1]
            else:
                form_elem = [func_id]
        else:
            if Src_Range[s][2] == func_id:
                form_elem = [func_id]
                func_str = ["curr", "volt"][func_id] + ["", ",time"][Common["time_stamp"]]
            else:
                form_elem = [0, 1]
                func_str = "volt,curr" + ["", ",time"][Common["time_stamp"]]
        if Common["time_stamp"] == 1:
            form_elem.append(3)
        s_obj.setformelem(form_elem)
        Common["var_name"][s] = func_str.split(",")

    for s in Meas_Range:
        s_obj = Common["smu_obj"][s]
        func_str = Func_Map[Meas_Range[s][2]]
        if Meas_Range[s][0] == 1:  # fixed range
            if srcrng_over20 == 1 and Meas_Range[s][2] == 0 and Meas_Range[s][1] >= 0.02:
                return reporterror("Measure range too big for 1KV source range.", MEAS_RNG_BIG)
            s_obj.range(Meas_Range[s][1], Meas_Range[s][2])
        elif Meas_Range[s][0] == 2:  # auto range with low limit
            s_obj.lorange(Meas_Range[s][1], Meas_Range[s][2])
        else:  # auto range
            s_obj.range(0, Meas_Range[s][2])
        # set plc
        s_obj.setnplc(Common["NPLC"], Meas_Range[s][2])

    for s in Remote_Sense:
        s_obj = Common["smu_obj"][s]
        if Remote_Sense[s] != -1:
            s_obj.setremotesense(Remote_Sense[s])
    for s in Open:
        s_obj = Common["smu_obj"][s]
        s_obj.setoffmode(0)
    # Calc Step value
    for s in Step:
        Step[s][0].append((Step[s][0][1] * 1.0 - Step[s][0][0]) / (Common["step_pts"] - 1))
    return 0


def reporterror(msg, code):
    print(code, msg)
    return code


def universalcom(cmd_dict, func_id=0):
    """
    func_id: 0="write",1="read",2="ask",3="ask for values"
    """
    rval = None
    if not debug_mode:
        print('universalcom', cmd_dict)
        for s in cmd_dict:
            print('s is:', s)
            cmd_str = ";".join(cmd_dict[s])
            print('cmd_str', cmd_str)
            for cmd_str in cmd_dict[s]:
                print("??/")
                if func_id == 0:
                    inst.write(s, cmd_str)
                elif func_id == 1:
                    rval = inst.read(s)
                elif func_id == 2:
                    rval = inst.ask(s, cmd_str)
                elif func_id == 3:
                    rval = inst.ask_for_values(s, cmd_str)
                elif func_id == 4:
                    rval = inst.read_values(s)
            cmd_dict[s] = []  # clear cmd dict
    else:
        for s in cmd_dict:
            for cmd in cmd_dict[s]:
                print("\n ", s, cmd)
                if func_id == 0:
                    inst.write(s, cmd)
                    if "?" not in cmd:
                        err = inst.ask(s, ":syst:err:all?")
                        print("error: ", err)
                elif func_id == 1:
                    rval = inst.read(s)
                elif func_id == 2:
                    rval = inst.ask(s, cmd)
                elif func_id == 3:
                    rval = inst.ask_for_values(s, cmd)
                elif func_id == 4:
                    rval = inst.read_values(s)
                print("\t", rval)
            cmd_dict[s] = []  # clear cmd dict
    print('rval is', rval)
    return rval
